
import numpy as np
from rank_bm25 import BM25Okapi

from src.spark_edu_rag.cache import redis_client
from src.spark_edu_rag.db import mysql_client
from src.spark_edu_rag.utils import get_logger, process_text


class BM25OkapiRetriever(object):
    def __init__(self,mysql_client,redis_client):
        self.logger = get_logger(__name__)
        self.mysql_client = mysql_client
        self.redis_client = redis_client

        self.bm25 = None
        self.question = None
        # 初始化问题列表，用于匹配答案对的原始文本对照
        self.original_questions = None
        self.data = self._load_data()

    def _load_data(self):
        original_key = "qa_original_questions"
        tokenized_key = "qa_tokenized_questions"

        # 从Redis缓存中获取原始问题列表
        self.original_questions = self.redis_client.get_str(original_key)
        # 从Redis缓存中获取已处理的问题列表
        tokenized_questions = self.redis_client.get_str(tokenized_key)
        if not self.original_questions or not tokenized_questions:
            # 如果缓存中不存在，从MySQL数据库中查询问题列表
            self.original_questions = self.mysql_client.fetch_questions()
            if not self.original_questions:
                self.logger.error("从MySQL数据库中查询问题列表失败")
                raise ValueError("从MySQL数据库中查询问题列表失败")

            # 对原始问题列表进行预处理，例如去除空格、特殊字符等
            tokenized_questions = [process_text(q[0]) for q in self.original_questions]
            # 存入Redis缓存中
            self.redis_client.add_str(original_key, [q[0] for q in self.original_questions])
            self.redis_client.add_str(tokenized_key, tokenized_questions)
        self.question = tokenized_questions
        self.bm25 = BM25Okapi(self.question)
        self.logger.info(f"BM25Okapi模型初始化完成，问题数量：{len(self.question)}")


    def _softmax(self, scores):
        """分数归一化，将BM25分数转换为概率分布， 用于后续阈值判断"""
        exp_scores = np.exp(scores - np.max(scores))
        return exp_scores / np.sum(exp_scores)

    def query(self,query, threshold=0.85):
        """根据输入查询问题，检索并返回相关问题"""
        # 1、对查询问题检查有效性，即为空，非字符串类型，抛出异常
        if not query or not isinstance(query, str):
            self.logger.error("无效查询问题")
            raise ValueError("无效查询问题")
        # 2、从Redis缓存中获取查询问题的答案
        cached_answer = self.redis_client.get_answer(query)
        if cached_answer:
            return cached_answer
        try:
            scores = self.bm25.get_scores(process_text(query))
            softmax_score = self._softmax(scores)
            best_index = softmax_score.argmax()
            best_score = softmax_score[best_index]
            if best_score >= threshold:
                result = self.original_questions[best_index]
                answer = self.mysql_client.fetch_answer(result[0])
                if answer:
                    # 缓存查询问题与答案对，过期时间为60秒
                    self.redis_client.add_str(query, answer[0])
                    self.logger.info(f"查询问题 {result[0]} 有对应的答案，相似度：{best_score:.4f}，缓存到Redis")
                    return answer[0]
                else:
                    self.logger.error(f"查询问题 {result[0]}， 相似度：{best_score:.4f}， 没有对应的答案")
                    return None
            else:
                return None


        except Exception as e:
            self.logger.error(f"检索问题,查询mysql时发生异常：{e}")
            raise ValueError(f"检索问题,查询mysql时发生异常：{e}")


bm25_search = BM25OkapiRetriever(mysql_client,redis_client)

if __name__ == '__main__':
    retriever = BM25OkapiRetriever(mysql_client,redis_client)

